file(
  GLOB TEST_OPS
  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
  "test_*.py")
string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")

list(REMOVE_ITEM TEST_OPS test_multi_precision_fp16_train)
list(REMOVE_ITEM TEST_OPS test_correlation)
foreach(src ${TEST_OPS})
  py_test(${src} SRCS ${src}.py)
endforeach()

py_test(test_correlation SRCS test_correlation.py ENVS "FLAGS_enable_pir_api=0")
py_test_modules(
  test_multi_precision_fp16_train
  MODULES
  test_multi_precision_fp16_train
  ENVS
  FLAGS_cudnn_deterministic=true
  FLAGS_cudnn_batchnorm_spatial_persistent=true
  FLAGS_conv_workspace_size_limit=1000)

set_tests_properties(test_multi_precision_fp16_train PROPERTIES TIMEOUT 120)
set_tests_properties(test_image_classification_fp16 PROPERTIES TIMEOUT 120)
